import networkx as nx
import pulp
import random
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D

# === Step 1: 随机生成燃气管网拓扑 ===
def generate_gas_network(num_nodes=20, num_edges=30, num_leaks=3):
    # 创建带权图
    G = nx.Graph()
    
    # 添加节点（编号0到n-1）
    G.add_nodes_from(range(num_nodes))
    
    # 添加边（管道）并随机分配管径、压力、长度
    for _ in range(num_edges):
        u, v = random.sample(range(num_nodes), 2)
        if not G.has_edge(u, v):
            G.add_edge(u, v, 
                       diameter=random.uniform(0.1, 0.5),  # 管径 (m)
                       pressure=random.uniform(3, 10),     # 压力 (bar)
                       length=random.uniform(50, 300))     # 管道长度 (m)
    
    # 添加用户数（每个节点代表一个用户点）
    for node in G.nodes:
        G.nodes[node]['users'] = random.randint(1, 10)  # 用户数
    
    # 随机选择泄漏点
    leak_points = random.sample(list(G.nodes), num_leaks)
    for node in leak_points:
        G.nodes[node]['is_leak'] = True
    
    return G, leak_points

# === Step 2: 最小割隔离泄漏点 ===
def min_cut_isolation(G, leak_points, source_node=0):
    # 假设所有泄漏点连接到超级源点，求最小割
    # 这里简化为对每个泄漏点单独求割，取并集
    valves_to_close = set()
    
    for leak in leak_points:
        if leak != source_node:  # 排除源点本身
            cut = nx.minimum_cut(G, source_node, leak)
            # 提取割边对应的阀门（假设每条边是一个阀门）
            for part in cut:
                for u, v in G.edges:
                    if (u in part and v not in part) or (v in part and u not in part):
                        valves_to_close.add((u, v))
    
    return valves_to_close

# === Step 3: 线性规划优化 ===
def optimize_closure(G, leak_points, valves_to_close):
    # 定义问题
    prob = pulp.LpProblem("Gas_Network_Valve_Closure", pulp.LpMinimize)
    
    # 变量：每个阀门是否关闭（0/1）
    valve_vars = {}
    for u, v in G.edges:
        valve_vars[(u, v)] = pulp.LpVariable(f"Valve_{u}_{v}", cat=pulp.LpBinary)
    
    # 目标函数：最小化停气用户数 + 泄漏总量 + 关阀时间
    # 参数权重（可根据实际情况调整）
    weight_users = 1.0
    weight_leak = 1.0
    weight_time = 1.0
    
    # 计算停气用户数：假设关闭阀门后，下游用户被隔离
    def get_isolated_users(valve_closure):
        # 简化模型：关闭阀门后，从源点出发可达的节点为未隔离
        G_copy = G.copy()
        for u, v in valve_closure:
            G_copy.remove_edge(u, v)
        reachable = nx.node_connected_component(G_copy, 0)
        isolated_users = sum(G.nodes[node]['users'] for node in G.nodes if node not in reachable)
        return isolated_users
    
    # 目标函数
    prob += (
        weight_users * pulp.lpSum(G.nodes[n]['users'] for n in G.nodes if n not in nx.node_connected_component(G, 0)),
        "Minimize_Isolated_Users"
    )
    prob += (
        weight_leak * pulp.lpSum(
            G[u][v]['diameter'] * G[u][v]['pressure'] * G[u][v]['length'] / 1000 for u, v in G.edges
        ),
        "Minimize_Leakage"
    )
    prob += (
        weight_time * pulp.lpSum(valve_vars[(u, v)] * random.uniform(0.5, 2.0) for u, v in G.edges),
        "Minimize_Closure_Time"
    )
    
    # 约束：必须关闭最小割集中的阀门
    for u, v in valves_to_close:
        prob += valve_vars[(u, v)] == 1, f"Must_Close_{u}_{v}"
    
    # 求解
    prob.solve()
    
    # 提取结果
    optimal_closures = [(u, v) for u, v in G.edges if pulp.value(valve_vars[(u, v)]) == 1]
    
    return optimal_closures

# === Step 4: 可视化 ===
def visualize_network(G, leak_points, valves_to_close, optimal_closures):
    pos = nx.spring_layout(G)
    plt.figure(figsize=(12, 8))
    
    # 绘制普通边
    nx.draw_networkx_nodes(G, pos, node_size=300, node_color='lightblue')
    nx.draw_networkx_labels(G, pos, font_size=10)
    
    # 绘制泄漏点
    nx.draw_networkx_nodes(G, pos, nodelist=leak_points, node_color='red', label='Leak Points')
    
    # 绘制普通边
    edges = [(u, v) for u, v in G.edges if (u, v) not in valves_to_close and (v, u) not in valves_to_close]
    nx.draw_networkx_edges(G, pos, edgelist=edges, edge_color='gray', width=1.5)
    
    # 绘制最小割边
    nx.draw_networkx_edges(G, pos, edgelist=valves_to_close, edge_color='orange', width=2.5, label='Min-Cut Valves')
    
    # 绘制优化后关闭的阀门
    nx.draw_networkx_edges(G, pos, edgelist=optimal_closures, edge_color='green', width=3, label='Optimal Closure')
    
    plt.legend(handles=[
        Line2D([0], [0], color='red', lw=4, label='Leak Points'),
        Line2D([0], [0], color='orange', lw=4, label='Min-Cut Valves'),
        Line2D([0], [0], color='green', lw=4, label='Optimal Closure'),
    ])
    plt.title("Gas Network Valve Closure Optimization")
    plt.show()

# === 主程序 ===
if __name__ == "__main__":
    # 生成拓扑
    G, leak_points = generate_gas_network(num_nodes=20, num_edges=30, num_leaks=3)
    print("Generated Network Nodes:", G.nodes)
    print("Leak Points:", leak_points)
    
    # 计算最小割
    valves_to_close = min_cut_isolation(G, leak_points, source_node=0)
    print("Valves to Close (Min-Cut):", valves_to_close)
    
    # 线性规划优化
    optimal_closures = optimize_closure(G, leak_points, valves_to_close)
    print("Optimal Closures:", optimal_closures)
    
    # 可视化
    visualize_network(G, leak_points, valves_to_close, optimal_closures)